import React, { useState, useRef, useEffect } from 'react';
import { Row, Col, Card, Button, Space, Tag, Divider, Checkbox, message } from 'antd';
import {
  SwapOutlined,
  CloseOutlined,
  CheckOutlined,
  SyncOutlined,
  CopyOutlined
} from '@ant-design/icons';
import { ModelType, MODEL_INFO } from '@/types/models';
import MessageBubble from '../MessageBubble';
import { useConversationStore } from '@/stores/conversation.store';
import { generationService, GenerationRequest } from '@/services/generation.service';
import './style.css';

interface ComparisonResult {
  model: ModelType;
  content: string;
  status: 'pending' | 'generating' | 'completed' | 'error';
  responseTime?: number;
  tokens?: number;
}

interface ModelComparisonViewProps {
  message: string;
  models?: ModelType[];
  onSelect?: (model: ModelType, content: string) => void;
  onClose?: () => void;
}

const ModelComparisonView: React.FC<ModelComparisonViewProps> = ({
  message,
  models = [],
  onSelect,
  onClose
}) => {
  const [results, setResults] = useState<Record<ModelType, ComparisonResult>>({});
  const [syncScroll, setSyncScroll] = useState(true);
  const scrollRefs = useRef<Record<string, HTMLDivElement | null>>({});
  const { comparisonModels, toggleComparisonModel } = useConversationStore();

  const activeModels = models.length > 0 ? models : comparisonModels;

  useEffect(() => {
    // Initialize results for active models
    const initialResults: Record<ModelType, ComparisonResult> = {};
    activeModels.forEach(model => {
      initialResults[model] = {
        model,
        content: '',
        status: 'pending'
      };
    });
    setResults(initialResults);

    // Start generation for each model
    activeModels.forEach(model => {
      generateWithModel(model, message);
    });
  }, [message, activeModels]);

  const generateWithModel = async (model: ModelType, prompt: string) => {
    const startTime = Date.now();
    const { modelParameters, currentPromptTemplate, promptVariables } = useConversationStore.getState();
    
    setResults(prev => ({
      ...prev,
      [model]: {
        ...prev[model],
        status: 'generating'
      }
    }));

    try {
      const generationRequest: GenerationRequest = {
        model,
        messages: [{ role: 'user', content: prompt }],
        parameters: modelParameters,
        stream: true,
        promptTemplate: currentPromptTemplate || undefined,
        promptVariables
      };

      // Use streaming generation for real-time updates
      let fullContent = '';
      try {
        for await (const chunk of generationService.streamGenerate(generationRequest)) {
          fullContent += chunk;
          setResults(prev => ({
            ...prev,
            [model]: {
              ...prev[model],
              content: fullContent,
              status: 'generating'
            }
          }));
        }
      } catch (streamError) {
        console.error(`Streaming error for ${model}:`, streamError);
        // Fallback to non-streaming
        const response = await generationService.generate({
          ...generationRequest,
          stream: false
        });
        
        if (response.error) {
          throw new Error(response.error);
        }
        
        fullContent = response.content;
      }

      const responseTime = Date.now() - startTime;

      setResults(prev => ({
        ...prev,
        [model]: {
          model,
          content: fullContent,
          status: 'completed',
          responseTime,
          tokens: fullContent.length // Rough estimation, should come from API response
        }
      }));
    } catch (error) {
      console.error(`Generation error for ${model}:`, error);
      setResults(prev => ({
        ...prev,
        [model]: {
          ...prev[model],
          status: 'error'
        }
      }));
    }
  };

  const handleScroll = (model: ModelType, event: React.UIEvent<HTMLDivElement>) => {
    if (!syncScroll) return;

    const source = event.currentTarget;
    const scrollPercentage = source.scrollTop / (source.scrollHeight - source.clientHeight);

    Object.entries(scrollRefs.current).forEach(([key, ref]) => {
      if (ref && key !== model) {
        const targetScrollTop = scrollPercentage * (ref.scrollHeight - ref.clientHeight);
        ref.scrollTop = targetScrollTop;
      }
    });
  };

  const handleSelect = (model: ModelType) => {
    const result = results[model];
    if (result && result.status === 'completed' && onSelect) {
      onSelect(model, result.content);
      const modelName = MODEL_INFO[model]?.name || model;
      message.success(`已选择 ${modelName} 的响应`);
    }
  };

  const handleCopy = (content: string) => {
    navigator.clipboard.writeText(content);
    message.success('已复制到剪贴板');
  };

  return (
    <div className="model-comparison-view">
      <div className="comparison-header">
        <Space>
          <SwapOutlined />
          <span className="comparison-title">模型对比模式</span>
          <Tag color="blue">{activeModels.length} 个模型</Tag>
        </Space>
        <Space>
          <Checkbox
            checked={syncScroll}
            onChange={(e) => setSyncScroll(e.target.checked)}
          >
            同步滚动
          </Checkbox>
          {onClose && (
            <Button
              type="text"
              icon={<CloseOutlined />}
              onClick={onClose}
            >
              关闭
            </Button>
          )}
        </Space>
      </div>

      <Divider />

      <div className="comparison-prompt">
        <strong>提示词：</strong>
        <span>{message}</span>
      </div>

      <Row gutter={16} className="comparison-results">
        {activeModels.map(model => {
          const result = results[model];
          const modelInfo = MODEL_INFO[model];
          
          // Skip if model info not found
          if (!modelInfo) {
            console.warn(`Model info not found for: ${model}`);
            return null;
          }
          
          return (
            <Col key={model} xs={24} md={12} lg={8}>
              <Card
                className="comparison-card"
                title={
                  <Space>
                    <span>{modelInfo?.name || model}</span>
                    {result?.status === 'generating' && (
                      <SyncOutlined spin />
                    )}
                    {result?.status === 'completed' && (
                      <CheckOutlined style={{ color: '#52c41a' }} />
                    )}
                  </Space>
                }
                extra={
                  <Space>
                    {result?.status === 'completed' && (
                      <>
                        <Button
                          type="text"
                          size="small"
                          icon={<CopyOutlined />}
                          onClick={() => handleCopy(result.content)}
                        />
                        <Button
                          type="primary"
                          size="small"
                          onClick={() => handleSelect(model)}
                        >
                          选择
                        </Button>
                      </>
                    )}
                  </Space>
                }
              >
                <div
                  ref={(ref) => { scrollRefs.current[model] = ref; }}
                  className="comparison-content"
                  onScroll={(e) => handleScroll(model, e)}
                >
                  {result?.status === 'pending' && (
                    <div className="comparison-pending">等待开始...</div>
                  )}
                  {result?.status === 'error' && (
                    <div className="comparison-error">生成失败，请重试</div>
                  )}
                  {(result?.status === 'generating' || result?.status === 'completed') && (
                    <div className="comparison-message">
                      <MessageBubble
                        message={{
                          id: `comparison-${model}`,
                          role: 'assistant',
                          content: result.content,
                          timestamp: new Date(),
                          status: result.status === 'generating' ? 'streaming' : 'sent'
                        }}
                        isStreaming={result.status === 'generating'}
                      />
                    </div>
                  )}
                </div>
                
                {result?.status === 'completed' && (
                  <div className="comparison-stats">
                    <Space split={<Divider type="vertical" />}>
                      <span>响应时间: {(result.responseTime! / 1000).toFixed(2)}s</span>
                      <span>Tokens: {result.tokens}</span>
                      <span>成本: ￥{((result.tokens || 0) * modelInfo.costPerToken).toFixed(4)}</span>
                    </Space>
                  </div>
                )}
              </Card>
            </Col>
          );
        })}
      </Row>
    </div>
  );
};

export default ModelComparisonView;